package shardedmap

import "sync"

type ShardedMap struct {
	shards []*shard
	shardN int
}

type shard struct {
	items map[string]interface{}
	lock  *sync.RWMutex
}

func NewShardedMap(shardN int) *ShardedMap {
	if shardN <= 0 {
		panic("shardN must be greater than 0")
	}
	sm := &ShardedMap{
		shardN: shardN,
	}
	sm.shards = make([]*shard, shardN)
	for i := 0; i < shardN; i++ {
		sm.shards[i] = &shard{
			items: make(map[string]interface{}),
			lock:  new(sync.RWMutex),
		}
	}
	return sm
}

func (sm *ShardedMap) Set(key string, value interface{}) {
	shardIndex := sm.getShardIndex(key)
	sm.shards[shardIndex].lock.Lock()
	sm.shards[shardIndex].items[key] = value
	sm.shards[shardIndex].lock.Unlock()
}

func (sm *ShardedMap) Get(key string) (interface{}, bool) {
	shardIndex := sm.getShardIndex(key)
	sm.shards[shardIndex].lock.RLock()
	value, ok := sm.shards[shardIndex].items[key]
	sm.shards[shardIndex].lock.RUnlock()
	return value, ok
}
func (sm *ShardedMap) ItemLen() int {
	itemLen := 0
	for _, shard := range sm.shards {
		shard.lock.RLock()
		itemLen += len(shard.items)
		shard.lock.RUnlock()
	}
	return itemLen
}

func (sm *ShardedMap) getShardIndex(key string) int {
	return int(fnv32(key)) % sm.shardN
}

func fnv32(key string) uint32 {
	var (
		hash  uint32 = 2166136261 // FNV offset basis 0x811c9dc5
		prime uint32 = 16777619   // FNV prime 0x01000193
	)
	for i := 0; i < len(key); i++ {
		hash = (hash * prime) ^ uint32(key[i])
	}
	return hash
}
